import os
from typing import Annotated

from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
from langchain_core.runnables.graph import MermaidDrawMethod
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI

from config.model_config import get_chat_openai_zhipu

class State(TypedDict):
    messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)

# 实际请以官方文档为准
llm = get_chat_openai_zhipu()

def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}

graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()

# # 生成图片
# graph_png = graph.get_graph().draw_mermaid_png()
# with open("graph.png", "wb") as f:
#     f.write(graph_png)


# try:
#     # display(Image(graph.get_graph().draw_mermaid_png()))
#     display(Image(graph.get_graph().draw_mermaid_png(draw_method=MermaidDrawMethod.PYPPETEER)))
# except Exception as e:
#     print('display error,', e)
#     # This requires some extra dependencies and is optional
#     pass


def stream_graph_updates(user_input: str):
    for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)

while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "一句话介绍langchain和LangGraph"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break
